import requests
import json
import time
from typing import Dict, List, Union, Callable
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from .base import BaseModel

class OpenAICompatibleModel(BaseModel):
    """OpenAI兼容的API模型实现"""
    
    def __init__(self, api_key: str, model_name: str, api_url: str):
        super().__init__(api_key, model_name)
        self.api_url = api_url
        self.supports_thinking = False
        self.supports_streaming = True  # 支持流式输出
        
        # 配置重试策略
        retry_strategy = Retry(
            total=3,  # 最大重试次数
            backoff_factor=2,  # 重试间隔 (2, 4, 8秒)
            status_forcelist=[429, 500, 502, 503, 504]  # 需要重试的HTTP状态码
        )
        
        # 创建会话并应用重试策略
        self.session = requests.Session()
        self.session.mount("https://", HTTPAdapter(max_retries=retry_strategy))
    
    def generate_response(self, messages: List[Dict[str, str]], enable_thinking: bool = False) -> str:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "messages": [{"role": msg["role"], "content": msg["content"]} for msg in messages],
            "max_tokens": 1000,
            "temperature": 0.7,
            "stream": False  # 明确设置为非流式
        }
        
        # 打印发送的JSON报文用于调试
        print("Sending request to API:")
        print(json.dumps(data, ensure_ascii=False, indent=2))
        
        try:
            response = self.session.post(self.api_url, headers=headers, json=data, timeout=60)
            response.raise_for_status()
            
            result = response.json()
            return result.get('choices', [{}])[0].get('message', {}).get('content', '').strip()
            
        except requests.exceptions.Timeout:
            raise Exception("连接超时，请检查网络状态或稍后重试")
        except requests.exceptions.ConnectionError:
            raise Exception("连接失败，请检查网络连接")
        except requests.exceptions.HTTPError as e:
            status_code = e.response.status_code if e.response is not None else None
            if status_code == 401:
                raise Exception("API密钥无效或已过期")
            elif status_code == 429:
                raise Exception("请求过于频繁，请稍后重试")
            elif status_code and status_code >= 500:
                raise Exception("服务器暂时不可用，请稍后重试")
            else:
                raise Exception(f"HTTP请求失败: {str(e)}")
        except requests.exceptions.RequestException as e:
            raise Exception(f"请求失败: {str(e)}")
        except Exception as e:
            raise Exception(f"生成回答时出错: {str(e)}")
    
    def generate_response_stream(self, messages: List[Dict[str, str]], 
                               callback: Callable[[str], None],
                               enable_thinking: bool = False) -> None:
        """流式生成响应"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "messages": [{"role": msg["role"], "content": msg["content"]} for msg in messages],
            "max_tokens": 1000,
            "temperature": 0.7,
            "stream": True  # 启用流式输出
        }
        
        try:
            response = self.session.post(self.api_url, headers=headers, json=data, timeout=60, stream=True)
            response.raise_for_status()
            
            # 处理流式响应
            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data: '):
                        data_line = line[6:]  # 去掉 'data: ' 前缀
                        if data_line == '[DONE]':
                            break
                        try:
                            chunk = json.loads(data_line)
                            delta = chunk.get('choices', [{}])[0].get('delta', {})
                            if 'content' in delta:
                                callback(delta['content'])
                        except json.JSONDecodeError:
                            continue
            
        except requests.exceptions.Timeout:
            raise Exception("连接超时，请检查网络状态或稍后重试")
        except requests.exceptions.ConnectionError:
            raise Exception("连接失败，请检查网络连接")
        except requests.exceptions.HTTPError as e:
            status_code = e.response.status_code if e.response is not None else None
            if status_code == 401:
                raise Exception("API密钥无效或已过期")
            elif status_code == 429:
                raise Exception("请求过于频繁，请稍后重试")
            elif status_code and status_code >= 500:
                raise Exception("服务器暂时不可用，请稍后重试")
            else:
                raise Exception(f"HTTP请求失败: {str(e)}")
        except requests.exceptions.RequestException as e:
            raise Exception(f"请求失败: {str(e)}")
        except Exception as e:
            raise Exception(f"生成回答时出错: {str(e)}")